import requests
import base64
import os
import subprocess
from datetime import datetime
import time

# 获取环境变量中的路径和 API Key
bozo = os.environ.get("BOZO") + "/Downloads/"
os.chdir(bozo)

def generate_image(prompt, api_key, image_url=None):
    """使用Gemini模型生成图像，支持纯文本提示或基于现有图像的修改"""
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    # 构建内容数组，至少包含文本提示
    content = [
        {
            "type": "text",
            "text": prompt
        }
    ]
    
    # 如果提供了图像URL，则添加到内容中（用于图像修改场景）
    if image_url:
        content.append({
            "type": "image_url",
            "image_url": {
                "url": image_url
            }
        })
    
    data = {
        "model": "google/gemini-2.5-flash-image-preview:free",
        "temperature": 0.5,
        "messages": [
            {
                "role": "user",
                "content": content  # 使用新的content结构
            }
        ]
    }
    
    try:
        response = requests.post(url, json=data, headers=headers)
        response.raise_for_status()
        result = response.json()
        
        # 提取图像URL
        image_url = result["choices"][0]["message"]["images"][0]["image_url"]["url"]
        return image_url
    except Exception as e:
        print(f"生成图像出错: {str(e)}")
        # 打印响应内容以便调试
        if response:
            print(f"响应内容: {response.text}")
        return None
    

def save_image_from_data_url(data_url, save_path):
    """从data URL保存图像到本地文件"""
    try:
        # 解析data URL
        if data_url.startswith('data:image/'):
            # 分割数据URL
            header, data = data_url.split(',', 1)
            # 解码Base64数据
            image_data = base64.b64decode(data)
            
            # 保存到文件
            with open(save_path, 'wb') as f:
                f.write(image_data)
            print(f"图像已保存到: {save_path}")
            return True
        else:
            print("无效的图像数据URL")
            return False
    except Exception as e:
        print(f"保存图像出错: {str(e)}")
        return False

def open_image_with_os(save_path):
    """使用系统默认程序打开图像文件（跨平台兼容）"""
    try:
        if os.name == 'nt':
            os.startfile(save_path)
        elif os.name == 'posix':
            subprocess.run(['open' if sys.platform == 'darwin' else 'xdg-open', save_path])
    except Exception as e:
        print(f"打开图像失败: {str(e)}")

def main():
    # 请替换为你的API密钥
    api_key = os.environ.get("ORAPI")
    
    # 输入提示词
    prompt = input("请输入修改的英文提示词：\r\n")
    image_url = input("请输入图片URL：\r\n")
    
    
    # 生成图像
    image_url = generate_image(prompt, api_key, image_url)
    if not image_url:
        print("无法生成图像，程序退出")
        return
    
    # 保存图像
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"OR_{timestamp}.png"
    save_path = os.path.join(os.getcwd(), filename)
    if not save_image_from_data_url(image_url, save_path):
        print("无法保存图像，程序退出")
        return
    
    # 打开图像
    # webbrowser.open(save_path)
    open_image_with_os(save_path)

if __name__ == "__main__":
    main()
